#!/usr/bin python3
# -*- coding: utf-8 -*-
# @Time    : 19-9-9 下午2:27
# @Author  : Aries
# @Site    : 
# @File    : get_mnist.py
# @Software: PyCharm

import os
from skimage import io
import torchvision.datasets.mnist as mnist

root="/home/ubuntu/Python-projects/Pytorch_test/fashion_mnist_data/FashionMNIST/raw"
save_path = "/home/ubuntu/Python-projects/Pytorch_test/fashion_mnist_data/FashionMNIST/Image_data/"
train_set = (
    mnist.read_image_file(os.path.join(root, 'train-images-idx3-ubyte')),
    mnist.read_label_file(os.path.join(root, 'train-labels-idx1-ubyte'))
        )
test_set = (
    mnist.read_image_file(os.path.join(root, 't10k-images-idx3-ubyte')),
    mnist.read_label_file(os.path.join(root, 't10k-labels-idx1-ubyte'))
        )
print("training set :",train_set[0].size())
print("test set :",test_set[0].size())
cls = ['Top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle_boot']

def convert_to_img(train=True):
    if(train):
        f=open(save_path+'train.txt','w')
        data_path=save_path+'train/'
        if(not os.path.exists(data_path)):
            os.makedirs(data_path)
        for i, (img,label) in enumerate(zip(train_set[0],train_set[1])):
            name_cls = cls[int(label)]
            if (not os.path.exists(data_path+name_cls)):
                os.makedirs(data_path+name_cls)
            img_path=data_path+name_cls+"/"+str(i)+'.jpg'
            io.imsave(img_path,img.numpy())
            f.write(img_path+'\n')
        f.close()
    else:
        f = open(save_path + 'test.txt', 'w')
        data_path = save_path + 'train/'
        if (not os.path.exists(data_path)):
            os.makedirs(data_path)
        for i, (img,label) in enumerate(zip(test_set[0],test_set[1])):
            name_cls = cls[int(label)]
            if (not os.path.exists(data_path + name_cls)):
                os.makedirs(data_path + name_cls)
            img_path = data_path + name_cls + "/" + str(i+60000) + '.jpg'
            io.imsave(img_path, img.numpy())

            f.write(img_path + '\n')
        f.close()

convert_to_img(True)
convert_to_img(False)